﻿//
// Created by QU on 24-9-28.
//

#include <boost/uuid/uuid_generators.hpp>

#include <iostream>
#include <print>
#include <utility>

#include "server/server/Session.h"
#include "server/logic/LogicNodeProcessor_Impl.cpp"

Session::Session(ip::tcp::socket&& socket, Worker& worker, std::stop_token stoken)
    : m_socket(std::move(socket))
      , m_uuid(boost::uuids::random_generator()())
      , m_worker(worker)
      , m_session_stop_token(std::move(stoken))
      , m_node_processor(std::make_shared<LogicNodeProcessor>())
{
    m_recv_buffer.resize(MIN_BUFFER_SIZE);

    std::cout << "Session created: " << getUuid() << std::endl;
}

Session::~Session()
{
    std::cout << "Session::~Session(): " << getUuid() << std::endl;
    close();
}

void Session::startSession()
{
    if (!m_is_started.exchange(true))
    {
        std::cout << "Session::startSession(): " << getUuid() << std::endl;
        auto self = shared_from_this();
        startAsyncRead();
    }
    else
    {
        std::cout << "Session already started: " << getUuid() << std::endl;
    }
}

bool Session::isStarted() const
{
    return m_is_started.load(std::memory_order_acquire);
}

bool Session::isClosed() const
{
    return !m_is_started.load(std::memory_order_acquire);
}

void Session::close()
{
    if (m_is_started.exchange(false))
    {
        std::cout << "Session::close(): closing session: " << getUuid() << std::endl;

        boost::system::error_code ec;

        m_socket.cancel(ec); // NOLINT(*-unused-return-value)
        if (ec)
        {
            std::cerr << "Error cancelling socket operations: " << ec.message() << std::endl;
        }

        m_socket.close(ec); // NOLINT(*-unused-return-value)
        if (ec)
        {
            std::cerr << "Error closing socket: " << ec.message() << std::endl;
        }

        // 清空发送队列
        {
            std::lock_guard lock(m_send_mutex);
            while (!m_send_queue.empty())
            {
                m_send_queue.pop();
            }
            m_current_send_node.reset();
        }

        m_recv_buffer.clear();
        m_recv_buffer_curr_len = 0;
        m_header_received_flag = false;
        m_temp_payload_size = 0;
        m_buffer_need_shrink_counter = 0;

        // 通知 Worker 移除此 Session
        std::cout << "Session on Worker closed: " << getUuid() << std::endl;
        m_worker.clearSession(getUuid());
    }
    else
    {
        std::cout << "Session already closed or not started: " << getUuid() << std::endl;
    }
}

void Session::startAsyncRead()
{
    if (m_session_stop_token.stop_requested())
    {
        close();
        return;
    }

    m_recv_buffer.resize(MIN_BUFFER_SIZE);
    m_recv_buffer_curr_len = 0;
    m_header_received_flag = false;
    m_temp_payload_size = 0;
    m_buffer_need_shrink_counter = 0;

    continueRead();
}

void Session::continueRead()
{
    if (m_session_stop_token.stop_requested())
    {
        close();
        return;
    }

    auto self = shared_from_this();
    m_socket.async_read_some(
        asio::buffer(m_recv_buffer.data() + m_recv_buffer_curr_len,
                     m_recv_buffer.size() - m_recv_buffer_curr_len),
        [this](const boost::system::error_code& ec, const std::size_t bytes_transferred)
        {
            if (m_session_stop_token.stop_requested())
            {
                close();
                return;
            }
            handleRead(ec, bytes_transferred);
        }
    );
}

void Session::handleRead(const boost::system::error_code& ec, std::size_t bytes_transferred)
{
    if (m_session_stop_token.stop_requested())
    {
        close();
        return;
    }

    if (ec)
    {
        handleError(ec);
        return;
    }

    std::print("Session::handleRead(): Received: {} bytes.\n", bytes_transferred);
    m_recv_buffer_curr_len += bytes_transferred;

    while (m_recv_buffer_curr_len > 0)
    {
        // 开始处理头部
        if (!m_header_received_flag)
        {
            // 如果收到的数据足够一个头
            if (m_recv_buffer_curr_len >= Node::MAX_HEADER_SIZE)
            {
                const uint16_t payload_size = boost::endian::big_to_native(
                    *reinterpret_cast<uint16_t*>(&m_recv_buffer[2]));

                std::print("Session::handleRead(): Header parse: payload_size: {} bytes.\n", payload_size);

                if (payload_size > Node::MAX_PAYLOAD_SIZE)
                {
                    std::cerr << "Invalid payload size: " << payload_size << std::endl;
                    handleError(ec, "Invalid payload size");
                    return;
                }

                // 这个包的头部, 存起来.
                m_temp_payload_size = payload_size;
                // 收到了这个包的头, 确定是否需要增加buffer的大小.
                if (const size_t required_size = m_temp_payload_size + Node::MAX_HEADER_SIZE; required_size >
                    m_recv_buffer.size())
                {
                    if (!adjustBufferSize(required_size))
                    {
                        // buffer扩容处理失败
                        handleError(ec, "Failed to adjust buffer size");
                        return;
                    }
                }

                // 标记头部处理完成
                m_header_received_flag = true;
            }
            else
            {
                break; // 需要更多数据来完成头部
            }
        }
        // 处理payload
        else
        {
            if (const size_t total_message_size = m_temp_payload_size + Node::MAX_HEADER_SIZE; m_recv_buffer_curr_len >=
                total_message_size)
            {
                const auto recv_node = nodeConstructor(m_recv_buffer, total_message_size);
                if (!recv_node)
                {
                    std::cerr << "Failed to process message." << std::endl;
                    handleError(ec, "Failed to process message");
                    return;
                }

                // 这里将shared_from_this传递到了NodeProcessor, 避免了Server中session map.clear, 直接导致Session被释放.
                processNode(*m_node_processor, shared_from_this(), recv_node);

                // 移动剩余数据到缓冲区开头
                const size_t remaining_data = m_recv_buffer_curr_len - total_message_size;
                if (remaining_data > 0)
                {
                    std::memmove(m_recv_buffer.data(), m_recv_buffer.data() + total_message_size, remaining_data);
                }
                m_recv_buffer_curr_len = remaining_data;

                // 重置状态
                m_header_received_flag = false;
                m_temp_payload_size = 0;
            }
            else
            {
                break; // 需要更多数据来完成消息
            }
        }
    }

    adjustBufferSize();
    continueRead();

    std::cout << "\t\t\t\t\t\t\t\t\t\tBuffer usage: " << m_recv_buffer_curr_len << "/" << m_recv_buffer.size() <<
        " bytes." << std::endl;
}

bool Session::adjustBufferSize(const size_t required_size)
{
    size_t new_size = m_recv_buffer.size();
    bool need_resize = false;

    if (required_size > 0 && required_size > new_size)
    {
        // 需要立即扩容buffer
        new_size = required_size;
        need_resize = true;
    }
    else if (m_recv_buffer_curr_len > static_cast<size_t>(EXPAND_THRESHOLD * new_size))
    {
        new_size = std::min(new_size * 2, MAX_BUFFER_SIZE);
        need_resize = new_size > m_recv_buffer.size();
    }
    else if (m_recv_buffer_curr_len < static_cast<size_t>(SHRINK_THRESHOLD * new_size))
    {
        if (++m_buffer_need_shrink_counter >= SHRINK_DELAY)
        {
            new_size = std::max(new_size / 2, MIN_BUFFER_SIZE);
            need_resize = new_size >= m_recv_buffer_curr_len * 2 && new_size < m_recv_buffer.size();
            if (need_resize)
            {
                m_buffer_need_shrink_counter = 0;
            }
        }
    }
    else
    {
        m_buffer_need_shrink_counter = 0; // 重置收缩计数器
    }

    if (need_resize)
    {
        try
        {
            m_recv_buffer.resize(new_size);
            std::cout << "\t\t\t\t\t\t\t\t\t\t[INFO]: Buffer " << (new_size > m_recv_buffer.size()
                                                                       ? "expanded"
                                                                       : "shrunk")
                << " to " << new_size << " bytes." << std::endl;
        }
        catch (const std::bad_alloc& e)
        {
            std::cerr << "Failed to resize buffer: " << e.what() << std::endl;
            return false;
        }
    }

    return true;
}

void Session::sendNode(const std::shared_ptr<Node>& node)
{
    bool start_write = false;
    {
        std::lock_guard lock(m_send_mutex);
        m_send_queue.push(node);
        if (!m_current_send_node)
        {
            start_write = true;
        }
    }

    if (start_write)
    {
        startAsyncWrite();
    }
}


void Session::startAsyncWrite()
{
    std::shared_ptr<Node> node_to_send;
    {
        std::lock_guard lock(m_send_mutex);
        if (!m_current_send_node && !m_send_queue.empty())
        {
            m_current_send_node = m_send_queue.front();
            m_send_queue.pop();
            node_to_send = m_current_send_node;
        }
    }

    if (node_to_send && !m_session_stop_token.stop_requested())
    {
        const auto [buffer, length] = node_to_send->serialize();

        auto self = shared_from_this();
        async_write(m_socket,
                    asio::buffer(buffer.get(), length),
                    [this](const boost::system::error_code& ec,
                           [[maybe_unused]] std::size_t bytes_transferred)
                    {
                        this->handleWrite(ec);
                    });
    }
    else if (m_session_stop_token.stop_requested())
    {
        close();
    }
}

void Session::handleWrite(const boost::system::error_code& ec)
{
    if (ec)
    {
        handleError(ec, "Failed to write data");
        return;
    }

    std::lock_guard lock(m_send_mutex);
    m_current_send_node.reset();

    if (m_session_stop_token.stop_requested() && m_send_queue.empty())
    {
        close();
    }

    else if (!m_send_queue.empty())
    {
        startAsyncWrite();
    }
}

void Session::handleError(const boost::system::error_code& ec, const std::string& message)
{
    if (!message.empty())
    {
        std::cerr << "Error: " << message << ": " << ec.message() << std::endl;
    }
    else
    {
        std::cerr << "Error on receive: " << ec.message() << std::endl;
    }
    close();
}

std::ostream& Session::printBuffer(std::ostream& os, const std::vector<std::byte>& buffer, const size_t buffer_size,
                                   const size_t header_size)
{
    constexpr size_t OUTPUT_HEAD_TAIL = 16;
    os << "Session::printBuffer()\n";
    os << "Buffer size: " << buffer_size << " Payload size: " << buffer_size - header_size << "\n";
    os << "Header (hex): ";
    os << std::hex << std::setfill('0');
    for (size_t i = 0; i < header_size; ++i)
    {
        if (i > 0) os << " ";
        os << std::setw(2) << static_cast<int>(buffer[i]);
    }
    os << "\n";

    os << "Payload (hex): ";
    for (size_t i = header_size; i < buffer_size && i < header_size + OUTPUT_HEAD_TAIL; ++i)
    {
        if (i > header_size) os << " ";
        os << std::setw(2) << static_cast<int>(buffer[i]);
    }
    if (buffer_size > header_size + OUTPUT_HEAD_TAIL) { os << " ..."; } // 如果buffer超过32字节，添加省略号
    for (size_t i = std::max(buffer_size - OUTPUT_HEAD_TAIL, header_size + OUTPUT_HEAD_TAIL); i < buffer_size; ++i)
    {
        if (i > header_size) os << " ";
        os << std::setw(2) << static_cast<int>(buffer[i]);
    }
    os << std::dec << "\n"; // 恢复为十进制输出

    return os;
}

std::shared_ptr<Node> Session::nodeConstructor(const std::vector<std::byte>& buffer, const size_t buffer_size)
{
    std::print("Session::nodeConstructor() Processed new message: \n");
    try
    {
        if (buffer_size < Node::MAX_HEADER_SIZE || buffer_size > buffer.size())
        {
            throw std::invalid_argument("Invalid buffer size");
        }

        auto node = std::make_shared<Node>(buffer.data(), static_cast<uint16_t>(buffer_size));

        // 在这里可以添加任何额外的处理逻辑
        printBuffer(std::cout, buffer, buffer_size);

        return node;
    }
    catch (const std::exception& e)
    {
        std::cerr << "Error processing buffer: " << e.what() << std::endl;
        return nullptr;
    }
}

void Session::processNode(NodeProcessor& processor, const std::shared_ptr<Session>& session,
                          const std::shared_ptr<Node>& node)
{
    if (!node || !session)
    {
        std::cerr << "Invalid node or session." << std::endl;
        return;
    }

    std::cout << "ProcessNode interface... " << std::endl;
    processor(session, node);
}

std::string Session::getUuid() const
{
    return to_string(m_uuid);
}
